import pandas as pd
from pathlib import Path

# 大字典，每个Kind每个Dist每一对相互作用
def All():
    AllPair = {}
    Kinds = ['HB', 'Arom', 'Phob', 'DC', 'other']
    Dists = ['S', 'M', 'L1', 'L2' ,'L3' ,'L4' ,'L5' ,'L6' ,'I1' ,'I2']
    AA    = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']

    for k in Kinds:
        AllPair[k] = {}
        for d in Dists:
            AllPair[k][d] = {}
            for i in range(len(AA)):
                a1 = AA[i]
                for j in range(i, len(AA)):
                    a2 = AA[j]
                    p = a1 + '-' + a2
                    AllPair[k][d][p] = []
    
    return AllPair


# 把每个Pair包含的具体相互作用装到一块
def Sort():
    AllPair = All()
    df = pd.read_csv('21D.csv')
    df.insert(1, 'Prop3', 'Undefined')
    Stable   = []
    Flexible = []
    Other    = []

    for i in range(len(df)):
        info = df.iloc[i].tolist()
        ac = [i, info[5], info[9], info[10], info[14], info[16], info[20], info[15], info[21]]
        AllPair[info[2]][info[3]][info[4]].append(ac)
        
    #    return AllPair


    #AllPair = Sort()
    for k in AllPair:
        for d in AllPair[k]:
            for p in AllPair[k][d]:
                pair = AllPair[k][d][p]
                if len(pair) >= 1000: # 起码要有1000个样本
                    As1 = [i[1] for i in pair]
                    Cs1 = [i[2] for i in pair]
                    As2 = [i[3] for i in pair]
                    Cs2 = [i[4] for i in pair]
                    Av  = [i[5] for i in pair]
                    Cv  = [i[6] for i in pair]
                    
                    # A是均值Ave，C是变异系数CV
                    As = As1 + As2
                    Cs = Cs1 + Cs2
                    As.sort()
                    Cs.sort()
                    Av.sort()
                    Cv.sort()
                    
                    # 均值和变异系数的最大值和最小值，从第5个开始，避免离谱值
                    Asm = As[-5]
                    Asn = As[0]
                    Avm = Av[-5]
                    Avn = Av[0]
                    Csm = Cs[-5]
                    Csn = Cs[0]
                    Cvm = Cv[-5]
                    Cvn = Cv[0]
                    
                    # 阈值，变量名的34567代表系数
                    As3 = (Asm - Asn) * 0.3 + Asn
                    As4 = (Asm - Asn) * 0.4 + Asn                    
                    As7 = (Asm - Asn) * 0.7 + Asn
                    As8 = (Asm - Asn) * 0.8 + Asn
                    Cs1 = (Csm - Csn) * 0.1 + Csn
                    Cs3 = (Csm - Csn) * 0.3 + Csn
                    Cs4 = (Csm - Csn) * 0.4 + Csn
                    Cs5 = (Csm - Csn) * 0.5 + Csn
                    
                    Av3 = (Avm - Avn) * 0.3 + Avn
                    Av4 = (Avm - Avn) * 0.4 + Avn                    
                    Av7 = (Avm - Avn) * 0.7 + Avn
                    Av8 = (Avm - Avn) * 0.8 + Avn
                    Cv1 = (Cvm - Cvn) * 0.1 + Cvn
                    Cv3 = (Cvm - Cvn) * 0.3 + Cvn
                    Cv4 = (Cvm - Cvn) * 0.4 + Cvn
                    Cv5 = (Cvm - Cvn) * 0.5 + Cvn
                    
                    for j in pair:
                        # 绝对稳定
                        if (As7 >= j[1] >= As4 and As7 >=j[3] >= As4) and (Av7 >= j[5] >= Av4) and (j[2] <= Cs1 and j[4] <= Cs1) and (j[6] <= Cv1) and (j[7] >= 0.9 and j[8] >= 0.9):
                            Stable.append(j[0])
                        elif (j[1] >= As8 and j[3] >= As8) and (j[5] >= Av8) and (j[7] >= 0.9 and j[8] >= 0.9):
                            Stable.append(j[0])
                        # 绝对柔性
                        elif (j[1] <= As3 and j[3] <= As3) and (j[5] <= Av3) and (j[2] >= Cs5 and j[4] >= Cs5) and (j[6] >= Cv5) and (j[7] <= 0.6 and j[8] <= 0.6):
                            Flexible.append(j[0])
                        # 绝对其它
                        elif (As7 >= j[1] >= As4 and As7 >=j[3] >= As4) and (Av7 >= j[5] >= Av4) and (((Cs4 >= j[2] >= Cs3 and Cs4 >= j[4] >= Cs3) and (Cs4 >= j[6] >= Cv3)) or (0.7 <= j[7] <= 0.8 and 0.7 <= j[8] <= 0.8)):
                            Other.append(j[0])
                        elif (j[1] >= As8 and j[3] >= As8) and (j[5] >= Av8) and (j[7] <= 0.6 and j[8] <= 0.6):
                            Other.append(j[0])
                        elif (j[1] <= As3 and j[3] <= As3) and (j[5] <= Av3) and (j[2] >= Cs5 and j[4] >= Cs5) and (j[6] >= Cv5) and (j[7] >= 0.9 and j[8] >= 0.9):
                            Other.append(j[0])
                        elif (j[1] <= As3 and j[3] <= As3) and (j[5] <= Av3) and (j[2] <= Cs1 and j[4] <= Cs1) and (j[6] <= Cv1):
                            Other.append(j[0])
                            
    # breakpoint()
    df.loc[Stable,   'Prop3'] = 'St'
    df.loc[Flexible, 'Prop3'] = 'Fl'
    df.loc[Other,    'Prop3'] = 'Ot'
    
    print('St: ', len(Stable))
    print('Fl: ', len(Flexible))
    print('Ot: ', len(Other))
    print('Total: ', len(df))
   
    return df


df1 = Sort().drop(labels=['PDB', 'Residue'], axis=1).drop_duplicates()    
df1.to_csv('22D_UR.csv', index=False)    
    
    